
#include <unitree_controller.h>
#include <unitree_cameras.h>

#include "unitree_legged_sdk/unitree_legged_sdk.h"
#include <math.h>
#include <iostream>
#include <unistd.h>
#include <string.h>

using namespace std;
using namespace cv;

float rate = 1.3;

Robot robot;
Cameras camera0(0);
Cameras camera1(1);
auto& cmd = robot.cmd;
auto& frame0 = camera0.frame;
auto& frame1 = camera1.frame;

float min_speed = 0.05f * rate;
float max_speed = 0.40f * rate;

int h_frame = 400;
int w_frame = 464;

int threshold_turn = h_frame / 4;
float p_side = 0.001;
float p_rect = 0.008;

int yellow_lower[3] = {15, 100, 120};
int yellow_upper[3] = {45, 255, 220};

int blue_lower[3] = {110,  80,  90};
int blue_upper[3] = {180, 240, 210};

cv::Rect box;



bool get_color_block(cv::Mat img, int color_lower[], int color_upper[], int square=0)
{
    Mat dst;
	Mat kernel;
	kernel = getStructuringElement(MORPH_RECT, Size(5, 5));
 
	std::vector<std::vector<Point>> contours;
	std::vector<Vec4i> hireachy;
	Point2f center;
    float radius = 20;

    Mat hsv;
	cvtColor(img, hsv, COLOR_BGR2HSV);
	inRange(hsv, 
            Scalar(color_lower[0], color_lower[1], color_lower[2]), 
            Scalar(color_upper[0], color_upper[1], color_upper[2]), 
            dst);
	//开操作
	morphologyEx(dst, dst, MORPH_OPEN, kernel);
	morphologyEx(dst, dst, MORPH_OPEN, kernel);
	morphologyEx(dst, dst, MORPH_OPEN, kernel);
	morphologyEx(dst, dst, MORPH_OPEN, kernel);
	//获取边界
	findContours(dst, contours, hireachy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE, Point(0, 0));
	//框选面积最大的边界
	if (contours.size() > 0)
	{
		double maxArea = 0;
		for (int i = 0; i < contours.size(); i++)
		{
			double area = contourArea(contours[static_cast<int>(i)]);
			if (area > maxArea)
			{
				maxArea = area;
				box = boundingRect(contours[static_cast<int>(i)]);
				minEnclosingCircle(contours[static_cast<int>(i)], center, radius);
			}
		}
        if(!square) return true;
        if(square && maxArea >= square) return true;
        else return false;
	}
    else return false;
}


void processing(cv::Mat frame)
{
    if(get_color_block(frame, blue_lower, blue_upper, 15000))
    {
        rectangle(frame, box, Scalar(255,0,0), 2);
    }
    if(get_color_block(frame, yellow_lower, yellow_upper))
    {
        rectangle(frame, box, Scalar(0,255,255), 2);
    }
}


void unload()
{
    camera0.save(processing, 111);
    int t = 260;
    for(int i = 0; i < t; i++)
    {
        usleep(10000);
        cmd.mode = 2;
        cmd.velocity[0] = max_speed / rate;
        robot.control();
    }

    t = 160;
    for(int i = 0; i < t; i++)
    {
        usleep(10000);
        cmd.mode = 1;
        cmd.euler[1] = -1.0f;
        robot.control();
    }

    t = 250;
    for(int i = 0; i < t; i++)
    {
        usleep(10000);
        cmd.mode = 2;
        cmd.velocity[0] = max_speed / rate;
        robot.control();
    }
}


void step()
{
    camera0.save(processing, 222);
    int t = 500;
    for(int i = 0; i < t; i++)
    {
        usleep(10000);
        cmd.mode = 2;
        cmd.footRaiseHeight = 0.2f;
        cmd.velocity[0] = max_speed / rate;
        cmd.velocity[1] = 0.02f;
        robot.control();
    }
}


void end()
{
    camera0.save(processing, 333);
    int t = (int)(30 / rate);
    for(int i = 0; i < t; i++)
    {
        usleep(10000);
        cmd.mode = 2;
        cmd.velocity[0] = max_speed;
        robot.control();
    }

    exit(0);
}


int main(int argc, char *argv[])
{
    int t = (int)(150 / rate);
    for(int i = 0; i < t; i++)
    {
        usleep(10000);
        cmd.mode = 2;
        cmd.velocity[0] = max_speed;
        robot.control();
    }

    int cout_turn = 0;
    int item = 0;  ///////////////////////////////////////////////
    while(true)
    {
        camera0.getframe();
        camera1.getframe();
        if(get_color_block(frame0, blue_lower, blue_upper, 9000))
        {
            item++;
            switch(item)
            {
                case 1:
                    cout << "dealing " << to_string(item) << endl;
                    unload();
                    cout << "finish " << to_string(item) << endl;
                    break;
                case 2:
                    cout << "dealing " << to_string(item) << endl;
                    step();
                    cout << "finish " << to_string(item) << endl;
                    break;
                case 3:
                    cout << "dealing " << to_string(item) << endl;
                    end();
                    cout << "finish " << to_string(item) << endl;
                    break;
                default:
                    exit(-1);
            }
        }
        else
        {
            // get_color_block(frame0, yellow_lower, yellow_upper);
            // if(box.tl().y >= threshold_turn)
            // {
            //     cout << "turning" << endl;
            //     cout_turn++;
            //     camera0.save(processing, cout_turn);
            //     t = 70;
            //     for(int i = 0; i < t; i++)
            //     {
            //         usleep(10000);
            //         cmd.mode = 2;
            //         cmd.yawSpeed = -2.4f;
            //         robot.control();
            //     }
            // }

            // if(!get_color_block(frame1, yellow_lower, yellow_upper, 100))
            // {
            //     t = 35;
            //     for(int i = 0; i < t; i++)
            //     {
            //         usleep(10000);
            //         cmd.mode = 2;
            //         cmd.velocity[0] = max_speed * 0.8;
            //         robot.control();
            //     }
            //     camera0.getframe();
            //     if(get_color_block(frame0, blue_lower, blue_upper, 5000))
            //     {
            //         cout << "test" << endl;
            //         continue;
            //     }

            //     cout << "turning" << endl;
            //     cout_turn++;
            //     camera0.save(processing, cout_turn*10+0);
            //     camera1.save(processing, cout_turn*10+1);
            //     t = 70;
            //     for(int i = 0; i < t; i++)
            //     {
            //         usleep(10000);
            //         cmd.mode = 2;
            //         cmd.yawSpeed = -2.4f;
            //         robot.control();
            //     }
            // }

            if(!get_color_block(frame1, yellow_lower, yellow_upper, 100))
            {
                cout << "turning" << endl;
                cout_turn++;
                camera1.save(processing, cout_turn);
                // camera1.save(processing, cout_turn*10+1);

                // t = (int)(30 / rate / rate / rate / rate);
                // for(int i = 0; i < t; i++)
                // {
                //     usleep(10000);
                //     cmd.mode = 2;
                //     cmd.velocity[0] = max_speed * 0.8;
                //     robot.control();
                // }
                t = 80;
                for(int i = 0; i < t; i++)
                {
                    usleep(10000);
                    cmd.mode = 2;
                    cmd.yawSpeed = -2.5f;
                    robot.control();
                }
            }

            get_color_block(frame0, yellow_lower, yellow_upper);
            int x_center = (box.br().x + box.tl().x) / 2;
            float error = (float)(w_frame / 2 - x_center);
            cmd.mode = 2;
            cmd.velocity[0] = max_speed;
            cmd.velocity[1] = p_side * error;
            cmd.yawSpeed = p_rect * error;
            robot.control();
        }

        usleep(10000);
    }


    return 0;
}
